#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@ File: runtime.py
@ Author: YuJX
@ Function : 计算软硬件环节设置相关

"""
import os
import numpy as np
from collections import defaultdict
import scipy.signal

from matplotlib import pyplot as plt
from torch.utils.tensorboard import SummaryWriter

class HistoryRecord():
    def __init__(self, log_dir, model, input_shape, record_type=['tb','txt']):
        self.record_type = record_type
        self.log_dir    = log_dir
        # self.epoch = []
        self.record = defaultdict(list)
        self.ckpt_record = defaultdict(list)
        self.bring_history = True

        # 建立保存记录的文件夹
        if not os.path.exists(self.log_dir):
            os.makedirs(self.log_dir)

        # 建立tensorboard
        self.writer     = SummaryWriter(self.log_dir)

        # TODO: 由于使用np而产生TraceWarning，目前关闭该功能
        # try:
        #     dummy_input     = torch.randn(1, 3, input_shape[0], input_shape[1])
        #     self.writer.add_graph(model, dummy_input)
        # except:
        #     pass

    def append_record (self,epoch,record):
        # train_val_record = {}
        if self.bring_history:
            self.bring_history = False

            for i in self.ckpt_record['epoch']:
                for key in self.ckpt_record.keys():
                    if key != 'epoch':
                        if 'tb' in self.record_type:
                            # 写入tensorboard
                            self.writer.add_scalar(key, self.ckpt_record[key][i-1], i)
                        if 'txt' in self.record_type:
                            # 写入.txt
                            with open(os.path.join(self.log_dir, key + ".txt"), 'a') as f:
                                f.write(str(self.ckpt_record[key][i-1]))
                                f.write("\n")

        self.ckpt_record['epoch'].append(epoch)

        for key in record.keys():
            self.ckpt_record[key].append(record[key])
            if 'tb' in self.record_type:
                # 写入tensorboard
                self.writer.add_scalar(key, record[key], epoch)
                self.writer.add_text(key, str(record[key]),epoch)
            if 'txt' in self.record_type:
                # 写入.txt
                with open(os.path.join(self.log_dir, key + ".txt"), 'a') as f:
                    f.write(str(record[key]))
                    f.write("\n")

        # 为每个指标分别绘图记录
        if 'plot' in self.record_type:
            for key in self.ckpt_record.keys():
                if key != 'epoch':
                    plt.plot(self.ckpt_record['epoch'], self.ckpt_record[key],'red', linewidth = 2,label=key)
                    plt.title(key)
                    plt.xlabel('epoch')
                    plt.ylabel(key)
                    plt.grid(True)
                    plt.legend(loc="upper right")
                    plt.savefig(os.path.join(self.log_dir, key + ".png"))
                    plt.cla()
                    plt.close("all")
            # 将train和val loss绘制在一张图中方便查看
            self.loss_plot()

    def loss_plot(self):
        plt.figure()
        plt.plot(self.ckpt_record['epoch'], self.ckpt_record['train_loss'], 'red', linewidth = 2, label='train loss')
        plt.plot(self.ckpt_record['epoch'], self.ckpt_record['val_loss'], 'coral', linewidth = 2, label='val loss')
        try:
            if len(self.ckpt_record['train_loss']) < 25:
                num = 5
            else:
                num = 15
            
            plt.plot(self.ckpt_record['epoch'], scipy.signal.savgol_filter(self.ckpt_record['train_loss'], num, 3), 'green', linestyle = '--', linewidth = 2, label='smooth train loss')
            plt.plot(self.ckpt_record['epoch'], scipy.signal.savgol_filter(self.ckpt_record['val_loss'], num, 3), '#8B4513', linestyle = '--', linewidth = 2, label='smooth val loss')
        except:
            pass

        plt.grid(True)
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.legend(loc="upper right")

        plt.savefig(os.path.join(self.log_dir, "train_val_loss.png"))

        plt.cla()
        plt.close("all")
def show_config(log_dir, **kwargs):
    print('Configurations:')
    print('-' * 80)
    print('|%25s | %50s|' % ('keys', 'values'))
    print('-' * 80)
    for key, value in kwargs.items():
        print('|%25s | %50s|' % (str(key), str(value)))
        with open(os.path.join(log_dir, "Mainconfig.txt"), 'a') as f:
            f.write('|%25s | %50s|' % (str(key), str(value)))
            f.write("\n")
    print('-' * 80)